"""
psql fleague work -q -c "select * from districts" >/tmp/districts.txt
"""
import sys
import os
import itertools

KEYWORD = "　"

filename = "最新县及县以上行政区划代码.txt"
querysql = 'psql fleague work -q -c "select code, id from districts" >/tmp/districts.txt'
importsql = 'psql fleague work -q -c "\i /tmp/district-{}.sql"'


def normalize(maxlevel):
    items = []
    
    with open(filename) as f:
        prelevel = 1
        precode = None
        stack = []
    
        for line in f:
            level = line.count(KEYWORD)
            code, district = line.strip().split()
    
            if level > prelevel:
                stack.append(precode)
            elif level < prelevel:
                for i in range(prelevel-level):
                    stack.pop()
    
            item = tuple(itertools.chain(stack, [code])), district
            items.append(item)
    
            prelevel = level
            precode = code
    
    items = list(filter(lambda i: len(i[0]) == maxlevel, items))
    
    os.system(querysql)

    with open('/tmp/districts.txt') as f:
        ids = [line.strip().split('|') for line in f]
        ids = {i[0].strip(): i[1].strip() for i in ids[2:-2]}

    # 直辖市修正
    fixed = [
        # 北京市
        ('110100', '110000'),
        ('110200', '110000'),
        # 天津市
        ('120100', '120000'),
        ('120200', '120000'),
        # 上海市
        ('310100', '310000'),
        ('310200', '310000'),
        # 重庆市
        ('500100', '500000'),
        ('500200', '500000'),
    ]

    for src, dest in fixed:
        if src in ids:
            ids[src] = ids[dest]

    #from pprint import pprint
    #pprint(items)

    sqls = []
    for item in items:
        code, district = item[0][-1], item[1]
        try:
            pcode = item[0][-2]
            pid = "'{}'".format(ids[pcode])
        except IndexError:  # maxlevel == 1
            pid = 'uuid_nil()'

        sql = "INSERT INTO districts VALUES (DEFAULT, {}, '{}', '{}', 1);".format(pid, district, code)
        sqls.append(sql)

    with open('/tmp/district-{}.sql'.format(maxlevel), 'w') as f:
        f.write("\n".join(sqls))

    os.system(importsql.format(maxlevel))


if __name__ == '__main__':
    normalize(1)
    normalize(2)
    normalize(3)
